#!/usr/bin/env python3
# set coding: utf-8
"""
@Project : litewaf2
@File :    models.py
@Time :    2023/2/18 14:45
@Author :  richard zhu
@Email :   gaotao.zhugt@dtzhejiang.com
@purpose : 
"""

from django.db import models
from portal.models import BaseModel
from django.utils.translation import gettext_lazy as _
from jsonschema import validate, SchemaError, ValidationError
import logging
import uuid
import re

logger = logging.getLogger(__name__)


rule_matchs_schema = {
    "title": "rule_matchs",
    "type": "array",
    "items": {
        "type": "object",
        "required": ["match_args", "args_prepocess", "match_operator", "match_value"],
        "properties": {
            "match_args": {
                "type": "array",
                "items": {
                    "type": "object",
                    "required": ["key", "value"],
                    "properties": {
                        "key": {
                            "type": "string"
                        },
                        "value": {
                            "type": "string"
                        }
                    },
                },
            },
            "args_prepocess": {
                "type": "array",
                "items": {
                    "type": "string"
                }
            },
            "match_operator": {
                "type": "string"
            },
            "match_value": {
                "type": "string"
            }
        },
    }
}


factor_keys_schema = {
    "title": "rule_matchs",
    "type": "object",
    "required": ["factor_keys", "factor_match"],
    "properties": {
        "factor_keys":{
            "type": "array",
            "items": {
                "type": "object",
                "required": ["key", "value"],
                "properties": {
                    "key": {
                        "type": "string"
                    },
                    "value": {
                        "type": "string"
                    }
                }
            }
        },
        "factor_match": {
            "type": "array",
            "items": {
                "type": "object",
                "required": ["action_value"],
                "properties": {
                    "action_value": {
                        "type": "object",
                        "required":["rate", "burst"],
                        "properties": {
                            "rate": {
                                "type": "string"
                            },
                            "burst": {
                                "type": "string"
                            }
                        }
                    }

                }
            }
        }
    }
}

class MyException(Exception):
    pass

class WafUser(BaseModel):
    user_name = models.CharField(max_length=120, unique=True)
    user_password = models.CharField(max_length=200)
    api_key = models.CharField(primary_key=True, auto_created=True,
                               default=re.split('-', str(uuid.uuid4()))[-1], max_length=100)
    api_password = models.CharField(auto_created=True, default=uuid.uuid4, max_length=100)

    def __unicode__(self):
        return str(self.api_key)

    def __str__(self):
        return "%s <%s>" % (self.user_name, self.api_key)


class GlobalSetting(BaseModel):
    name = models.CharField(primary_key=True, max_length=128, verbose_name=_("配置名"), unique=True)
    # uuid = models.CharField(primary_key=True, auto_created=True, default=uuid.uuid4,
    #                         editable=False, max_length=100)
    default_page = models.JSONField(verbose_name=_("serverName默认页"), default=dict)
    redis = models.JSONField(verbose_name=_("redis配置"), default=dict)
    sys_abnormal_handle_data = models.JSONField(verbose_name=_("失败处理"), default=dict)
    waf_module_enable = models.JSONField(verbose_name=_("打开模块"), default=dict)
    auto_update_period = models.IntegerField(default=30)
    flow_protection = models.JSONField(verbose_name=_("流量保护"), default=dict)
    block_page = models.JSONField(verbose_name=_("阻断返回信息"), default=dict)
    bot_check_srv_url = models.CharField(max_length=128, verbose_name=_("验证码服务URL"), blank=True, null=True)
    def __unicode__(self):
        return self.name

    def __str__(self):
        return "%s" % (self.name)



class FlowRule(BaseModel):
    RULE_TYPE = (
        ("flow_protection_rule", "流量保护规则"),
        ("flow_limit_burst_rule", "限流规则")
    )
    RULE_ACTION = (
        ("burst_limit_redis", "基于参数的redis限流"),
        ("burst_limit_local", "基于参数的本地限流"),
        ("flow_protection_local", "流量保护本地限流"),
        ("flow_protection_redis", "流量保护redis限流"),
    )
    rule_name = models.CharField(max_length=64, verbose_name=_("规则名称"))
    rule_type = models.CharField(max_length=64, choices=RULE_TYPE, verbose_name="规则类型")
    action_value = models.JSONField(verbose_name="限流参数", default=dict, blank=True, null=True)
    rule_action = models.CharField(max_length=64, choices=RULE_ACTION,verbose_name="处理规则")
    rule_matchs = models.JSONField(verbose_name="匹配规则", default=dict)
    uuid = models.CharField(default=uuid.uuid4, auto_created=True, unique=True, max_length=100)
    flow_factor_keys = models.JSONField(verbose_name="限流因子", default=dict, blank=True, null=True)
    order = models.IntegerField(default=5)


    def __unicode__(self):
        return self.rule_name

    def __str__(self):
        return "%s" % (self.rule_name)

    def save(self, *args, **kwargs):
        try:
            validate(instance=self.rule_matchs, schema=rule_matchs_schema)
        except SchemaError as e:
            print("验证模式schema出错：\n出错位置：{}\n提示信息：{}".format(" --> ".join([i for i in e.path]), e.message))
        except ValidationError as e:
            print("json数据不符合schema规定：\n出错字段：{}\n提示信息：{}".format(" --> ".join([i for i in e.path]), e.message))
        else:
            print("验证成功！")
            super().save(*args, **kwargs)
            return

        raise MyException("json格式验证错误")

class WafRule(BaseModel):
    RULE_ACTION = (
        ("block", "阻断"),
        ("allow", "允许"),
        ("reject_response", "拒绝"),
        ("bot_check", "验证码"),
        ("bypass", "观察模式")
    )
    rule_name = models.CharField(max_length=64, verbose_name=_("规则名称"))
    rule_type = models.CharField(max_length=64, default="waf_rule", auto_created=True, editable=False)
    rule_action = models.CharField(max_length=64, choices=RULE_ACTION, verbose_name="处理规则")
    rule_matchs = models.JSONField(verbose_name="匹配规则", default=dict)
    uuid = models.CharField(default=uuid.uuid4, auto_created=True, unique=True, max_length=100)
    order = models.IntegerField(default=5)

    def __unicode__(self):
        return self.rule_name

    def __str__(self):
        return "%s" % (self.rule_name)

    def save(self, *args, **kwargs):
        try:
            validate(instance=self.rule_matchs, schema=rule_matchs_schema)
        except SchemaError as e:
            print("验证模式schema出错：\n出错位置：{}\n提示信息：{}".format(" --> ".join([i for i in e.path]), e.message))
        except ValidationError as e:
            print("json数据不符合schema规定：\n出错字段：{}\n提示信息：{}".format(" --> ".join([i for i in e.path]), e.message))
        else:
            print("验证成功！")
            super().save(*args, **kwargs)
            return

        raise MyException("json格式验证错误")


class RuleGroup(BaseModel):
    name = models.CharField(max_length=128)
    waf_rule = models.ManyToManyField(WafRule, blank=True)
    flow_rule = models.ManyToManyField(FlowRule, blank=True)

    def waf_rule_list(self):
        return ','.join([a.rule_name for a in self.waf_rule.all()])

    def flow_rule_list(self):
        return ','.join([a.rule_name for a in self.flow_rule.all()])

    def __unicode__(self):
        return self.name

    def __str__(self):
        return "%s" % (self.name)

class WafDomain(BaseModel):
    user_id = models.ForeignKey(WafUser, on_delete=models.DO_NOTHING)
    domain = models.CharField(max_length=500)
    port = models.IntegerField(default=80)
    rulegroup = models.ManyToManyField(RuleGroup, blank=True)

    def rg_list(self):
        return ','.join([a.name for a in self.rulegroup.all()])

    def __unicode__(self):
        return self.domain

    def __str__(self):
        return "%s" % (self.domain)

class Node(BaseModel):
    user_id = models.ForeignKey(WafUser, on_delete=models.DO_NOTHING)
    hostname = models.CharField(max_length=300)
    uuid = models.CharField(primary_key=True, unique=True, max_length=200)
    ip = models.GenericIPAddressField()
    content = models.TextField(verbose_name="主机备注", blank=True, null=True)

    def __str__(self):
        return "%s" % (self.hostname)

    def __unicode__(self):
        return self.hostname